# !/usr/bin/env python3
# coding=utf8
"""
主要根据"成交数据"和"每日收盘价数据"做统计分析,
"""
import collections
import datetime
import numpy
import pandas
from typing import Dict, List, Set, Tuple, Optional, Union, Callable
from vnpy.tools.bar_generator import BarGeneratorVnpy
from vnpy.trader.object import TradeData
from vnpy_ctastrategy.backtesting import DailyResult


class StatisticsAnalysis(object):
    """
    源自 vnpy_ctastrategy/vnpy_ctastrategy/backtesting.py 的 BacktestingEngine
    """

    def __init__(self) -> None:
        """"""
        self.rate: float = 0  # 手续费率
        self.slippage: float = 0  # 交易滑点
        self.size: float = 1  # 合约乘数
        self.capital: int = 1_000_000  # 回测资金
        self.risk_free: float = 0
        self.annual_days: int = 240

        self.trades: Dict[str, TradeData] = {}
        self.daily_results: Dict[datetime.date, DailyResult] = {}
        self.statistics: dict = None

    def set_parameters(
        self,
        trades: List[TradeData],
        prices: Dict[datetime.datetime, float],
        rate: float,
        slippage: float,
        size: float,
        capital: int = 0,
        risk_free: float = 0,
        annual_days: int = 240
    ) -> None:
        """
        prices: 为了和"vnpy_ctastrategy/backtesting.py 的 BacktestingEngine"逻辑保持相似,
        需求, 字典的key是datetime.datetime(要求date部分是business_day), 不是datetime.date,
        """
        self.trades = {trade.vt_tradeid: trade for trade in trades}
        self.update_daily_close(prices=prices)

        self.rate = rate
        self.slippage = slippage
        self.size = size
        self.capital = capital

        self.risk_free = risk_free
        self.annual_days = annual_days

    def execute(self, output: bool) -> None:
        """"""
        daily_df: Optional[pandas.DataFrame] = self.calculate_result()
        self.statistics: dict = self.calculate_statistics(daily_df, output=output)

    def update_daily_close(self, prices: Dict[datetime.datetime, float]) -> None:
        """"""
        self.daily_results.clear()

        for self_datetime, price in prices.items():

            d: datetime.date = self_datetime.date()

            daily_result: Optional[DailyResult] = self.daily_results.get(d, None)
            if daily_result:
                daily_result.close_price = price
            else:
                self.daily_results[d] = DailyResult(d, price)

    def calculate_result(self) -> Optional[pandas.DataFrame]:
        """"""
        if not self.trades:
            return None

        # Add trade data into daily reuslt.
        for trade in self.trades.values():
            d: datetime.date = BarGeneratorVnpy.business_day(trade)  # d: datetime.date = trade.datetime.date()
            daily_result: DailyResult = self.daily_results[d]
            daily_result.add_trade(trade)

        # Calculate daily result by iteration.
        pre_close = 0
        start_pos = 0

        for daily_result in self.daily_results.values():
            daily_result.calculate_pnl(
                pre_close,
                start_pos,
                self.size,
                self.rate,
                self.slippage
            )

            pre_close = daily_result.close_price
            start_pos = daily_result.end_pos

        # Generate dataframe
        results: collections.defaultdict = collections.defaultdict(list)

        for daily_result in self.daily_results.values():
            for key, value in daily_result.__dict__.items():
                if key == "trades":  # 这个字段在DataFrame里面没有用到,
                    continue
                results[key].append(value)

        daily_df = pandas.DataFrame.from_dict(results).set_index("date")

        return daily_df

    def calculate_statistics(self, df: pandas.DataFrame = None, output=True) -> dict:
        """"""
        # Check for init DataFrame
        if df is None:
            # Set all statistics to 0 if no trade.
            start_date: str = ""
            end_date: str = ""
            total_days: int = 0
            profit_days: int = 0
            loss_days: int = 0
            end_balance: float = 0
            max_drawdown: float = 0
            max_ddpercent: float = 0
            max_drawdown_duration: int = 0
            total_net_pnl: float = 0
            daily_net_pnl: float = 0
            total_commission: float = 0
            daily_commission: float = 0
            total_slippage: float = 0
            daily_slippage: float = 0
            total_turnover: float = 0
            daily_turnover: float = 0
            total_trade_count: int = 0
            daily_trade_count: int = 0
            total_return: float = 0
            annual_return: float = 0
            daily_return: float = 0
            return_std: float = 0
            sharpe_ratio: float = 0
            return_drawdown_ratio: float = 0
        else:
            # Calculate balance related time series data
            df["balance"] = df["net_pnl"].cumsum() + self.capital

            # When balance falls below 0, set daily return to 0
            pre_balance: pandas.Series = df["balance"].shift(1)
            pre_balance.iloc[0] = self.capital
            x = df["balance"] / pre_balance
            x[x <= 0] = numpy.nan
            df["return"] = numpy.log(x).fillna(0)

            df["highlevel"] = (
                df["balance"].rolling(
                    min_periods=1, window=len(df), center=False).max()
            )
            df["drawdown"] = df["balance"] - df["highlevel"]
            df["ddpercent"] = df["drawdown"] / df["highlevel"] * 100

            # Calculate statistics value
            start_date = df.index[0]
            end_date = df.index[-1]

            total_days: int = len(df)
            profit_days: int = len(df[df["net_pnl"] > 0])
            loss_days: int = len(df[df["net_pnl"] < 0])

            end_balance = df["balance"].iloc[-1]
            max_drawdown = df["drawdown"].min()
            max_ddpercent = df["ddpercent"].min()
            max_drawdown_end = df["drawdown"].idxmin()

            if isinstance(max_drawdown_end, datetime.date):
                max_drawdown_start = df["balance"][:max_drawdown_end].idxmax()
                max_drawdown_duration: int = (max_drawdown_end - max_drawdown_start).days
            else:
                max_drawdown_duration: int = 0

            total_net_pnl: float = df["net_pnl"].sum()
            daily_net_pnl: float = total_net_pnl / total_days

            total_commission: float = df["commission"].sum()
            daily_commission: float = total_commission / total_days

            total_slippage: float = df["slippage"].sum()
            daily_slippage: float = total_slippage / total_days

            total_turnover: float = df["turnover"].sum()
            daily_turnover: float = total_turnover / total_days

            total_trade_count: int = df["trade_count"].sum()
            daily_trade_count: int = total_trade_count / total_days

            total_return: float = (end_balance / self.capital - 1) * 100
            annual_return: float = total_return / total_days * self.annual_days
            daily_return: float = df["return"].mean() * 100
            return_std: float = df["return"].std() * 100

            if return_std:
                daily_risk_free: float = self.risk_free / numpy.sqrt(self.annual_days)
                sharpe_ratio: float = (daily_return - daily_risk_free) / return_std * numpy.sqrt(self.annual_days)
            else:
                sharpe_ratio: float = 0

            if max_ddpercent:
                return_drawdown_ratio: float = -total_return / max_ddpercent
            else:
                return_drawdown_ratio = 0

        # Output
        if output:
            self.output("-" * 30)
            self.output(f"首个交易日：\t{start_date}")
            self.output(f"最后交易日：\t{end_date}")

            self.output(f"总交易日：\t{total_days}")
            self.output(f"盈利交易日：\t{profit_days}")
            self.output(f"亏损交易日：\t{loss_days}")

            self.output(f"起始资金：\t{self.capital:,.2f}")
            self.output(f"结束资金：\t{end_balance:,.2f}")

            self.output(f"总收益率：\t{total_return:,.2f}%")
            self.output(f"年化收益：\t{annual_return:,.2f}%")
            self.output(f"最大回撤: \t{max_drawdown:,.2f}")
            self.output(f"百分比最大回撤: {max_ddpercent:,.2f}%")
            self.output(f"最长回撤天数: \t{max_drawdown_duration}")

            self.output(f"总盈亏：\t{total_net_pnl:,.2f}")
            self.output(f"总手续费：\t{total_commission:,.2f}")
            self.output(f"总滑点：\t{total_slippage:,.2f}")
            self.output(f"总成交金额：\t{total_turnover:,.2f}")
            self.output(f"总成交笔数：\t{total_trade_count}")

            self.output(f"日均盈亏：\t{daily_net_pnl:,.2f}")
            self.output(f"日均手续费：\t{daily_commission:,.2f}")
            self.output(f"日均滑点：\t{daily_slippage:,.2f}")
            self.output(f"日均成交金额：\t{daily_turnover:,.2f}")
            self.output(f"日均成交笔数：\t{daily_trade_count}")

            self.output(f"日均收益率：\t{daily_return:,.2f}%")
            self.output(f"收益标准差：\t{return_std:,.2f}%")
            self.output(f"Sharpe Ratio：\t{sharpe_ratio:,.2f}")
            self.output(f"收益回撤比：\t{return_drawdown_ratio:,.2f}")

        statistics: dict = {
            "start_date": start_date,
            "end_date": end_date,
            "total_days": total_days,
            "profit_days": profit_days,
            "loss_days": loss_days,
            "capital": self.capital,
            "end_balance": end_balance,
            "max_drawdown": max_drawdown,
            "max_ddpercent": max_ddpercent,
            "max_drawdown_duration": max_drawdown_duration,
            "total_net_pnl": total_net_pnl,
            "daily_net_pnl": daily_net_pnl,
            "total_commission": total_commission,
            "daily_commission": daily_commission,
            "total_slippage": total_slippage,
            "daily_slippage": daily_slippage,
            "total_turnover": total_turnover,
            "daily_turnover": daily_turnover,
            "total_trade_count": total_trade_count,
            "daily_trade_count": daily_trade_count,
            "total_return": total_return,
            "annual_return": annual_return,
            "daily_return": daily_return,
            "return_std": return_std,
            "sharpe_ratio": sharpe_ratio,
            "return_drawdown_ratio": return_drawdown_ratio,
        }

        # Filter potential error infinite value
        for key, value in statistics.items():
            if value in (numpy.inf, -numpy.inf):
                value = 0
            statistics[key] = numpy.nan_to_num(value)

        return statistics

    def output(self, msg) -> None:
        """
        Output message of backtesting engine.
        """
        print(f"{datetime.datetime.now()}\t{msg}")


if __name__ == "__main__":
    pass
